# -*- coding:utf-8 -*-
"""
@Time : 2022/4/24 10:15 AM
@Author: binkuolo
@Des: mysql数据库
"""
from typing import Type
from fastapi import FastAPI
from tortoise.contrib.fastapi import register_tortoise
from tortoise.models import Model

import os


class Router:
    def db_for_read(self, model: Type[Model]):
        return "slave"

    def db_for_write(self, model: Type[Model]):
        return "master"


# -----------------------数据库配置-----------------------------------
DB_ORM_CONFIG = {
    "connections": {
        "master": {
            'engine': 'tortoise.backends.mysql',
            "credentials": {
                'host': os.getenv('DB_HOST', '127.0.0.1'),
                'user': os.getenv('DB_USER', 'root'),
                'password': os.getenv('DB_PASSWORD', '123123'),
                'port': int(os.getenv('DB_PORT', 3306)),
                'database': os.getenv('DB_NAME', 'test'),
            }
        },
        "slave": {
            'engine': 'tortoise.backends.mysql',
            "credentials": {
                'host': os.getenv('SLAVE_DB_HOST', '127.0.0.1'),
                'user': os.getenv('SLAVE_DB_USER', 'root'),
                'password': os.getenv('SLAVE_DB_PASSWORD', '123123'),
                'port': int(os.getenv('SLAVE_DB_PORT', 3306)),
                'database': os.getenv('SLAVE_DB_NAME', 'test'),
            }
        },

    },
    "apps": {
        "base": {"models": ["models.base"], "default_connection": "master"},
    },
    # "routers": ["database.mysql.Router"],
    'use_tz': False,
    'timezone': 'Asia/Shanghai'
}

if os.getenv('RW_SEPARATE', '0') == 1:
    DB_ORM_CONFIG['routers'] = ["database.mysql.Router"]


async def register_mysql(app: FastAPI):
    # 注册数据库
    return register_tortoise(
        app,
        config=DB_ORM_CONFIG,
        generate_schemas=False,
        add_exception_handlers=False,
    )
